/*
给定一个含有 M x N 个元素的矩阵（M 行，N 列），请以对角线遍历的顺序返回这个矩阵中的所有元素，对角线遍历如下图所示。

 

示例:

输入:
[
 [ 1, 2, 3 ],
 [ 4, 5, 6 ],
 [ 7, 8, 9 ]
]

输出:  [1,2,4,7,5,3,6,8,9]

解释:

 

说明:

给定矩阵中的元素总数不会超过 100000 。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/diagonal-traverse
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
*/

#include "../stdc++.h"

// 对角线迭代和反转
class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        if (mat.empty() || mat[0].empty()) {
            return {};
        }
        int m = mat.size();
        int n = mat[0].size();
        vector<int> res(m * n, 0);
        int k{0};
        for (int d{0}; d < m + n - 1; ++d) {
            vector<int> temp{};
            int r = d < n ? 0 : d - n + 1; // 起始点的行坐标
            int c = d < n ? d : n - 1;     // 起始点的列坐标
            while (r < m && c > -1) {
                temp.push_back(mat[r][c]);
                ++r;
                --c;
            }
            // 反转
            if (d % 2 == 0) {
                reverse(temp.begin(), temp.end());
            }
            for (auto& num : temp) {
                res[k++] = num;
            }
        }
        return res;
    }
};

// 模拟
class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        if (mat.empty() || mat[0].empty()) {
            return {};
        }
        int m = mat.size();
        int n = mat[0].size();
        vector<int> res(m * n, 0);
        int index{0};
        int row{0};
        int col{0};
        int direction{1};
        while (row < m && col < n) {
            res[index++] = mat[row][col];
            int new_row = row + (direction == 1 ? -1 : 1);
            int new_col = col + (direction == 1 ? 1 : -1);
            if (new_col < 0 || new_col == n || new_row < 0 || new_row == m) {
                if (direction == 1) {
                    row += (col == n - 1 ?  1 : 0);
                    col += (col < n - 1 ? 1 : 0);
                } else {
                    col += (row == m - 1 ? 1 : 0);
                    row += (row < m - 1 ? 1 : 0);
                }
                direction = 1 - direction;
            } else {
                row = new_row;
                col = new_col;
            }
        }
        return res;
    }
};